package org.seedframework.sharding.mybatis;


import com.google.gson.Gson;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.seedframework.sharding.annotations.Scatter;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import java.lang.reflect.Field;
import java.util.List;
import java.util.Set;


/**
 * Dynamic locate table interceptor
 *
 * <p>Identity table name by {@link Scatter}</p>
 */
@Slf4j
@Intercepts(
        {
                @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
                @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        }
)
public class ExecutorInterceptor implements Interceptor {


    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        String methodName = invocation.getMethod().getName();
        if (methodName.equals("query")) {
            return doQuery(invocation);
        } else if (methodName.equals("update")) {
            Executor executor = (Executor) invocation.getTarget();
            if (!"org.apache.ibatis.executor.CachingExecutor".equals(executor.getClass().getName())) {
                return doUpdate(invocation);
            }
        }
        return invocation.proceed();
    }


    private Object doQuery(Invocation invocation) throws Throwable {
        Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameter = args[1];
        Executor executor = (Executor) invocation.getTarget();

        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler resultHandler = (ResultHandler) args[3];
        CacheKey cacheKey;
        BoundSql boundSql;

        if (args.length == 4) {
            boundSql = mappedStatement.getBoundSql(parameter);
            cacheKey = executor.createCacheKey(mappedStatement, parameter, rowBounds, boundSql);
        } else {
            cacheKey = (CacheKey) args[4];
            boundSql = (BoundSql) args[5];
        }
        // original sql
        String querySql = boundSql.getSql();
        String newSql = querySql;
        // mapper method full path
        String id = mappedStatement.getId();
        ScatterRoutingDefinition definition = ScatterRoutingSupport.getMethodScatterStrategy(id);
        if (definition != null) {
            newSql = triggerNewSql(querySql, parameter, definition, "query");
            String dataSource = definition.dataSource();
            executor = ExecutorFactory.getExecutor(dataSource, executor);
        }
        boundSql.setSql(newSql);
        Object rst = executor.query(mappedStatement, parameter, rowBounds, resultHandler, cacheKey, boundSql);
        boundSql.setSql(querySql);
        return rst;
    }

    private Object doUpdate(Invocation invocation) throws Throwable {
        Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameter = args[1];
        Executor executor = (Executor) invocation.getTarget();
        BoundSql boundSql = mappedStatement.getBoundSql(parameter);
        String updateSql = boundSql.getSql();
        log.debug("==> Original SQL: {}", updateSql);
        String newSql = updateSql;
        // mapper method full path
        String id = mappedStatement.getId();
        String groupName = MultiWriteSupport.getGroupName(id);
        Object rst = null;
        ScatterRoutingDefinition definition = ScatterRoutingSupport.getMethodScatterStrategy(id);
        if (StringUtils.hasText(groupName)) {
            // do multi-datasource write, when do this, do not change sql, just switch datasource
            Set<Executor> executorSet = ExecutorFactory.getExecutors(groupName,executor);
            for(Executor exe1 : executorSet){
                rst = exe1.update(mappedStatement,parameter);
            }
        } else {
            if (definition != null) {
                newSql = triggerNewSql(updateSql, parameter, definition, "update");
                String dataSource = definition.dataSource();
                executor = ExecutorFactory.getExecutor(dataSource, executor);
            }
            log.debug("==> SQL EXEC: {}    PARAMETER: {}", newSql, new Gson().toJson(parameter));
            boundSql.setSql(newSql);
            mappedStatement.setBoundSql(boundSql);
            invocation.setTarget(executor);
            rst = invocation.proceed();
        }
        boundSql.setSql(updateSql);
        return rst;

    }


    private String triggerNewSql(String sql, Object parameter, ScatterRoutingDefinition scatterStrategy, String type) throws Throwable {

        MapperMethod.ParamMap paramMap = (MapperMethod.ParamMap) parameter;
        String scatterValue;
        String[] routes = scatterStrategy.getFieldsRoute();
        if (routes.length == 0) {
            scatterValue = (String) paramMap.get(scatterStrategy.getKey());
        } else {
            Object fObj = null;
            for (int i = 0; i < routes.length; i++) {
                if (i == 0) {
                    fObj = paramMap.get(routes[0]);
                } else {
                    fObj = getFieldValueByObject(fObj, routes[i]);
                }
            }
            scatterValue = (String) getFieldValueByObject(fObj, scatterStrategy.getKey());
        }

        scatterStrategy.setValue(scatterValue);
        String tableSuffix = scatterStrategy.tableSuffix();
        // replace table name in sql
        Statement statement = CCJSqlParserUtil.parse(sql);
        TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
        List<String> tableNames = tablesNamesFinder.getTableList(statement);
        Assert.notEmpty(tableNames, "No table name found in SQL.");
        String tableName = tableNames.get(0);
        sql = sql.replaceAll(tableName, tableName + tableSuffix);
        return sql;
    }


    private Object getFieldValueByObject(Object object, String targetFieldName) throws Exception {

        Class objClass = object.getClass();
        Field[] fields = objClass.getDeclaredFields();
        for (Field field : fields) {
            String currentFieldName = field.getName();
            if (currentFieldName.equals(targetFieldName)) {
                field.setAccessible(true);
                return field.get(object);
            }
        }
        return null;
    }



}
